from openai import OpenAI
from typing import List, Dict, Optional
from config import OPENAI_CONFIG, BILI_CONFIG
import re
import requests

def format_response(response: str) -> str:
    """
    格式化AI回复文本，使其更适合TTS使用
    
    Args:
        response: AI的原始回复文本
        
    Returns:
        格式化后的文本
    """
    # 1. 去除多余的空格和换行
    text = ' '.join(response.split())
    
    # 2. 去除重复的标点符号
    text = re.sub(r'[。！？.!?]+(?=[。！？.!?])', '', text)
    
    # 3. 规范化标点符号
    # 将英文标点转换为中文标点
    punctuation_map = {
        ',': '，',
        '.': '。',
        '!': '！',
        '?': '？',
        ':': '：',
        ';': '；'
    }
    for en, cn in punctuation_map.items():
        text = text.replace(en, cn)
    
    # 4. 去除代码块标记和其他markdown标记
    text = re.sub(r'```[\s\S]*?```', '', text)
    text = re.sub(r'`.*?`', '', text)
    text = re.sub(r'\*\*?(.*?)\*\*?', r'\1', text)
    
    # 5. 处理括号内容
    # 移除英文括号及其内容
    text = re.sub(r'\(.*?\)', '', text)
    # 保留中文括号内容，但去除括号本身
    text = re.sub(r'（(.*?)）', r'\1', text)
    
    # 6. 确保句子以合适的标点结尾
    if not text[-1] in ['。', '！', '？']:
        text += '。'
    
    return text.strip()

class BiliHandler:
    def __init__(
        self,
        sessdata: str = BILI_CONFIG["SESSDATA"],
        bili_jct: str = BILI_CONFIG["BILI_JCT"],
        buvid3: str = BILI_CONFIG["BUVID3"]
    ):
        """
        初始化BiliHandler
        
        Args:
            sessdata: B站的SESSDATA cookie
            bili_jct: B站的bili_jct cookie
            buvid3: B站的BUVID3 cookie
        """
        self.cookies = {
            "SESSDATA": sessdata,
            "bili_jct": bili_jct,
            "buvid3": buvid3
        }
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
            "Referer": "https://www.bilibili.com",
            "Origin": "https://www.bilibili.com"
        }
        self._verify_login()

    def _verify_login(self) -> bool:
        """
        验证登录状态
        
        Returns:
            bool: 是否成功登录
        """
        url = "https://api.bilibili.com/x/web-interface/nav"
        response = requests.get(url, cookies=self.cookies, headers=self.headers)
        data = response.json()
        
        if data["code"] != 0:
            raise ValueError("B站登录验证失败，请检查cookie设置")
        return True

class AIHandler:
    def __init__(
        self,
        api_key: str = OPENAI_CONFIG.get("api_key", ""),
        base_url: str = OPENAI_CONFIG.get("base_url", "https://api.openai.com/v1"),
        model: str = OPENAI_CONFIG.get("model", "gpt-3.5-turbo"),
        initial_prompt: str = "请你扮演我的助手，与我聊天。",
        debug: bool = OPENAI_CONFIG.get("debug", False),
        default_response: str = OPENAI_CONFIG.get("default_response", "对不起，我现在不能理解你的话。")
    ):
        """
        初始化AIHandler
        
        Args:
            api_key: OpenAI API密钥
            base_url: API基础URL地址
            model: 使用的模型名称
            initial_prompt: 初始系统提示语
            debug: 是否启用调试模式
            default_response: 错误时的默认回复
        """
        if not api_key:
            raise ValueError("OpenAI API密钥不能为空")
            
        self.client = OpenAI(
            api_key=api_key,
            base_url=base_url
        )
        self.model = model
        self.initial_prompt = initial_prompt
        self.debug = debug
        self.default_response = default_response
        self.messages: List[Dict[str, str]] = [
            {"role": "system", "content": initial_prompt}
        ]

    def reset_prompt(self, new_prompt: str) -> None:
        """
        重置系统提示语
        
        Args:
            new_prompt: 新的系统提示语
        """
        self.initial_prompt = new_prompt
        self.messages = [
            {"role": "system", "content": new_prompt}
        ]

    def chat(self, user_input: str) -> str:
        """
        与AI进行对话
        
        Args:
            user_input: 用户输入的消息
            
        Returns:
            AI的回复消息
        """
        try:
            # 添加用户消息到历史记录
            self.messages.append({"role": "user", "content": user_input})
            
            # 调用OpenAI API获取回复
            response = self.client.chat.completions.create(
                model=self.model,
                messages=self.messages
            )
            
            # 获取AI回复
            ai_response = response.choices[0].message.content
            
            # 将AI回复添加到历史记录
            self.messages.append({"role": "assistant", "content": ai_response})
            
            # 返回格式化后的回复
            return format_response(ai_response)
            
        except Exception as e:
            if self.debug:
                raise e
            else:
                # 非调试模式下返回默认回复
                # 将默认回复添加到历史记录
                self.messages.append({"role": "assistant", "content": self.default_response})
                return self.default_response

    def flush(self) -> None:
        """
        清除对话历史，只保留初始prompt
        """
        self.messages = [
            {"role": "system", "content": self.initial_prompt}
        ]
